from backend.api import (
    add_task,
    cancel_task,
    get_can_do_tasks,
    read_exp_info,
    read_exp_result,
)
from datetime import datetime
import json
import pandas as pd
import streamlit as st
from trade.data.table_api import ms_to_datetime_no_tz


if __name__ == "__main__":

    st.title("任务中心")

    st.subheader("可添加任务")

    with st.form("convert_date"):
        left, center, right = st.columns([3, 2, 3])
        with left:
            din = st.text_input(
                label="输入日期", key="input_date", label_visibility="collapsed"
            )
        with right:
            try:
                ms_out = int(datetime.strptime(din, "%Y%m%d%H%M%S").timestamp() * 1000)
                st.text(ms_out)
            except Exception as e:
                try:
                    ms_out = int(datetime.fromisoformat(din).timestamp() * 1000)
                    st.text(ms_out)
                except Exception as e:
                    st.text(e)
        with center:
            st.form_submit_button("转化日期")
    with st.form("convert_ms"):
        left, center, right = st.columns([3, 2, 3])
        with left:
            din = st.text_input(
                label="输入毫秒", key="input_ms", label_visibility="collapsed"
            )
        with right:
            try:
                st.text(ms_to_datetime_no_tz(int(din)).isoformat())
            except Exception as e:
                st.text(e)
        with center:
            st.form_submit_button("转化毫秒")
    can_do_tasks = get_can_do_tasks()
    for i, x in enumerate(can_do_tasks):
        with st.expander(x["task_name"]):
            with st.form(key=f"expand_form_{i}"):
                title = st.text_input("任务标题", key=f"expand_title_{i}")
                params = st.text_area(
                    "任务参数",
                    key=f"expand_params_{i}",
                    value=json.dumps(x.get("template_params"), indent=2),
                )
                if st.form_submit_button("提交任务"):
                    tid = add_task(x["task_name"], title, json.loads(params))
                    if tid:
                        st.success(f"添加任务成功: {tid}")
                    else:
                        st.error(f"添加任务失败")

    st.subheader("任务列表")
    if st.button("刷新任务表"):
        read_exp_info.clear()
        st.success("刷新成功")

    tasks = read_exp_info()
    tasks.sort(key=lambda x: -x["create"])
    for i in range(len(tasks)):
        tasks[i]["info_str"] = json.dumps(tasks[i]["info"])
    df = pd.DataFrame(tasks)
    for col in ["create", "start", "end"]:
        df[col] = pd.to_datetime(df[col], unit="ms")
    # df.sort_values("create", inplace=True, ascending=False)
    choose = st.dataframe(
        df[["expid", "name", "title", "create", "start", "end", "status", "info_str"]],
        on_select="rerun",
        selection_mode="single-row",
    )
    if choose.selection.rows:
        info = tasks[choose.selection.rows[0]]
        st.subheader("选择的任务")

        if info["status"] in ("RUNNING", "WAIT") and st.button("终止任务"):
            rs = cancel_task(info["expid"])
            if rs:
                st.success(f"取消任务成功: {info['expid']}")
            else:
                st.error(f"取消任务失败: {info['expid']}")

        elif info["status"] in ("FINISHED",):
            with st.expander("展开结果"):
                result = read_exp_result(info["expid"])
                print("运行结果是", info["expid"], result)
                if result is not None:
                    st.json(result)
                else:
                    st.error(f"获取结果失败: {info['expid']}")

        with st.expander("展开信息"):
            st.json(info)
        with st.expander("复制修改"):
            with st.form(key="modify_from"):
                title = st.text_input("任务标题")
                # print(json.dumps(info["info"]["kwargs"]))
                params = st.text_area(
                    "任务参数", value=json.dumps(info["info"]["kwargs"], indent=2)
                )
                print("任务参数look")
                if st.form_submit_button("提交任务"):
                    print("要提交任务")
                    tid = add_task(info["name"], title, json.loads(params))
                    if tid:
                        st.success(f"添加任务成功: {tid}")
                        print(f"复制添加成功: {tid}")
                    else:
                        st.error(f"添加任务失败")
                        print(f"复制添加失败: {tid}")
